{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/data02/users/lz/miniconda3/envs/UICoder/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import shutil\n",
    "from datasets import Dataset, concatenate_datasets\n",
    "from tqdm import tqdm\n",
    "\n",
    "path = '/data02/users/lz/code/UICoder/datasets/c4-wash/H128-2560_C128-4096_R2/c4-format-marked/origin'\n",
    "output_path = '/data02/users/lz/code/UICoder/datasets/c4-wash/H128-2560_C128-4096_R2/c4-format-marked/merged'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "194191"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# len(os.listdir('/data02/users/lz/code/UICoder/datasets/c4-wash/H128-2560_C128-4096_R2/c4-format'))\n",
    "# os.path.dirname('/data02/users/lz/code/UICoder/datasets/c4-wash/H128-2560_C128-4096_R2/c4-format')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['1k_0', '1k_1', '1k_2', '1k_3', '6k_5', '6k_6']"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# index = 5\n",
    "# start = 20000\n",
    "# size = 3000\n",
    "\n",
    "# source_path = '/data02/users/lz/code/UICoder/datasets/c4-wash/H128-2560_C128-4096_R2/c4-format'\n",
    "# target_path = f'{os.path.dirname(source_path)}/c4-format-{int(size/1000)}k_{index}'\n",
    "# for item in os.listdir(source_path)[start:start+size]:\n",
    "#     shutil.copytree(os.path.join(source_path,item),os.path.join(target_path,item))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1k_0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating train split: 0 examples [00:00, ? examples/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1000/1000 [01:38<00:00, 10.11it/s]s/ examples]\n",
      "Generating train split: 1000 examples [01:38, 10.11 examples/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1k_1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1000/1000 [01:46<00:00,  9.43it/s]s/ examples]\n",
      "Generating train split: 1000 examples [01:46,  9.43 examples/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1k_2\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1000/1000 [01:41<00:00,  9.87it/s]s/ examples]\n",
      "Generating train split: 1000 examples [01:41,  9.87 examples/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1k_3\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1000/1000 [01:43<00:00,  9.67it/s]s/ examples]\n",
      "Generating train split: 1000 examples [01:43,  9.67 examples/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "6k_5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3000/3000 [06:52<00:00,  7.27it/s]s/ examples]\n",
      "Generating train split: 3000 examples [06:52,  7.27 examples/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "6k_6\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3000/3000 [06:33<00:00,  7.62it/s]s/ examples]\n",
      "Generating train split: 3000 examples [06:33,  7.62 examples/s]\n"
     ]
    }
   ],
   "source": [
    "ds_avrs = []\n",
    "for batch_name in sorted(os.listdir(path)):\n",
    "    print(batch_name)\n",
    "    dss = []\n",
    "    for one in os.listdir(os.path.join(path,batch_name)):\n",
    "        ds = Dataset.load_from_disk(os.path.join(path,batch_name,one))\n",
    "        dss.append(ds)\n",
    "\n",
    "    def generator():\n",
    "        for idx,item in tqdm(enumerate(dss[0]),total=len(dss[0])):\n",
    "            image = item['image']\n",
    "            scores = []\n",
    "            item2 = {\n",
    "                'image': image,\n",
    "            }\n",
    "            for idx2,ds in enumerate(dss):\n",
    "                score = ds[idx]['score']\n",
    "                item2[f'score_{idx2}'] = score\n",
    "                scores.append(score)\n",
    "            item2['score_avr'] = round(sum(scores)/len(scores))\n",
    "            yield item2\n",
    "            \n",
    "\n",
    "    ds_avr = Dataset.from_generator(generator,)\n",
    "    ds_avrs.append(ds_avr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_merged = concatenate_datasets(ds_avrs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving the dataset (3/3 shards): 100%|██████████| 10000/10000 [00:03<00:00, 2930.28 examples/s]\n"
     ]
    }
   ],
   "source": [
    "ds_merged.save_to_disk(output_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "dsx = Dataset.load_from_disk(output_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10000"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(dsx)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
